import os
import yaml
from collections import defaultdict
from tqdm import tqdm


class DatasetClassCounter:
    def __init__(self, dataset_root, data_yaml_path, subsets_to_scan=None):
        self.dataset_root = dataset_root
        self.data_yaml_path = data_yaml_path
        if subsets_to_scan is None:
            self.subsets_to_scan = ['train', 'val', 'test']
        else:
            self.subsets_to_scan = subsets_to_scan
        self.class_names = {}
        self.subset_class_counts = defaultdict(lambda: defaultdict(int))
        self.total_class_counts = defaultdict(int)
        self.subset_image_counts = defaultdict(int)
    def _load_class_names(self):
        try:
            with open(self.data_yaml_path, 'r', encoding='utf-8') as f:
                data_yaml = yaml.safe_load(f)
            
            names_list = data_yaml.get('names', [])
            if isinstance(names_list, dict):
                for idx, name in names_list.items():
                    self.class_names[int(idx)] = name
            elif isinstance(names_list, list):
                for idx, name in enumerate(names_list):
                    self.class_names[idx] = name
            
            if not self.class_names:
                print(f"警告: {self.data_yaml_path} 中未找到 'names' 字段或为空。将仅显示类别ID。")
            else:
                print(f"从 {self.data_yaml_path} 读取到 {len(self.class_names)} 个类别。")
                for idx in sorted(self.class_names.keys()):
                    print(f"  ID {idx}: {self.class_names[idx]}")

        except FileNotFoundError:
            print(f"错误: 找不到 {self.data_yaml_path} 文件。请检查路径并确保文件存在。")
            exit()
        except yaml.YAMLError as e:
            print(f"错误: 解析 {self.data_yaml_path} 失败: {e}")
            exit()

    def _get_subsets_to_scan(self):
        test_labels_path = os.path.join(self.dataset_root, 'test', 'labels')
        test_images_path = os.path.join(self.dataset_root, 'test', 'images')
        
        if (os.path.exists(test_labels_path) or 
            (os.path.exists(test_images_path) and os.path.exists(test_labels_path))):
            if 'test' not in self.subsets_to_scan:
                self.subsets_to_scan.append('test')

    def _scan_subset(self, subset_name):
        labels_dir = os.path.join(self.dataset_root, subset_name, 'labels')
        images_dir = os.path.join(self.dataset_root, subset_name, 'images')

        if not self._check_labels_dir(labels_dir, subset_name):
            return

        print(f"\n正在扫描 {subset_name} 标签文件...")
        
        # 统计指定子目录下的图片数量
        self._count_images(images_dir, subset_name)
        label_files = self._get_label_files(labels_dir)

        # 依次处理每个标签文件，统计类别实例数量
        for label_filename in tqdm(label_files, desc=f"处理 {subset_name} 标签"):
            label_path = os.path.join(labels_dir, label_filename)
            self._process_label_file(label_path, label_filename, subset_name)

    def _check_labels_dir(self, labels_dir, subset_name):
        if not os.path.exists(labels_dir):
            print(f"警告: 找不到 {subset_name} 标签目录: {labels_dir}，跳过此子集。")
            return False
        return True

    def _count_images(self, images_dir, subset_name):
        if os.path.exists(images_dir):
            self.subset_image_counts[subset_name] = len([f for f in os.listdir(images_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))])

    def _get_label_files(self, labels_dir):
        return [f for f in os.listdir(labels_dir) if f.endswith('.txt')]

    def _process_label_file(self, label_path, label_filename, subset_name):
        try:
            with open(label_path, 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if parts:
                        try:
                            class_id = int(parts[0])
                            self.subset_class_counts[subset_name][class_id] += 1
                            self.total_class_counts[class_id] += 1
                        except ValueError:
                            print(f"警告: 文件 {label_filename} 包含无效的类别ID: '{parts[0]}'")
                            continue
        except Exception as e:
            print(f"错误: 读取标签文件 {label_path} 失败: {e}")

    def _print_report(self):
        print("\n--- 数据集类别分布报告 ---")

        print("\n图片数量统计:")
        total_overall_images = 0
        for subset in self.subsets_to_scan:
            count = self.subset_image_counts[subset]
            print(f"  {subset.capitalize()} 图片总数: {count}")
            total_overall_images += count
        print(f"总图片数: {total_overall_images}")

        print("\n类别实例数量统计 (按目录):")
        for subset_name in self.subsets_to_scan:
            print(f"\n目录: /{subset_name.capitalize()}")
            current_subset_total = sum(self.subset_class_counts[subset_name].values())
            if current_subset_total == 0:
                print("  无类别实例。")
                continue
            
            for class_id in sorted(self.subset_class_counts[subset_name].keys()):
                count = self.subset_class_counts[subset_name][class_id]
                class_name = self.class_names.get(class_id, f"未知类别 {class_id}")
                percentage = (count / current_subset_total) * 100 if current_subset_total > 0 else 0
                print(f"  - {class_name} (ID {class_id}): {count} 个实例 ({percentage:.2f}%)")

        print("\n总体类别实例数量统计:")
        overall_total_instances = sum(self.total_class_counts.values())
        if overall_total_instances == 0:
            print("整个数据集无类别实例。")
        else:
            for class_id in sorted(self.total_class_counts.keys()):
                count = self.total_class_counts[class_id]
                class_name = self.class_names.get(class_id, f"未知类别 {class_id}")
                percentage = (count / overall_total_instances) * 100
                print(f"  - {class_name} (ID {class_id}): {count} 个实例 ({percentage:.2f}%)")

    def run(self):
        class Colors:
            BLUE = '\033[92m'
            RESET = '\033[0m'
        print(f"\n\n{Colors.BLUE} *** 欢迎使用数据集类别数量分析工具 作者: marss ***{Colors.RESET}")
        print(f"\n--- 正在分析数据集目录<{self.dataset_root}> ---")

        # step1: 加载类别名称和子集信息
        self._load_class_names()
        self._get_subsets_to_scan()

        # step2: 扫描数据集子集
        for subset_name in self.subsets_to_scan:
            self._scan_subset(subset_name)
        
        # step3: 打印统计结果
        self._print_report()
        print("\n分析完成。")


if __name__ == "__main__":
    counter = DatasetClassCounter(
        dataset_root='datasets_balanced',
        data_yaml_path='hhf.yaml',
        subsets_to_scan=['train', 'valid', 'test']
    )
    counter.run()